import csv
import tokenize
import zipfile
import os

from .codeql_grammer_obj import *

def astnode_create(csv_row:dict):
    node_type=csv_row["node_type"]
    node_value=csv_row["node_value"]
    scope=csv_row["scope"]
    filepath=csv_row["filepath"]
    location=csv_row["location"]
    isLeaf=csv_row["isLeaf"]

    scope=Scope.parse(scope)
    location=NodeLocation.parse(filepath,location)
    
    return AstNode(node_type=node_type,
        is_leaves_node=isLeaf,
        value=node_value,
        scope=scope,
        location=location)
        
def astedge_create(csv_row:dict):
    parent_type=csv_row["start_type"]
    child_type=csv_row["end_type"]
    parent_value=csv_row["start_value"]
    child_value=csv_row["end_value"]
    
    filepath=csv_row["filepath"]
    child_location=csv_row["end_location"]
    parent_location=csv_row["start_location"]
    
    child_location_obj=NodeLocation.parse(filepath,child_location)
    parent_location_obj=NodeLocation.parse(filepath,parent_location)
    
    return AstEdge(AstNode.id_construct(parent_type,parent_value,parent_location_obj),
                   AstNode.id_construct(child_type,child_value,child_location_obj),
                   filepath)
    
def cfgnode_create(csv_row:dict):
    node_type=csv_row["node_type"]
    node_value=csv_row["node_value"]
    filepath=csv_row["filepath"]
    node_location=csv_row["node_location"]
    block_location=csv_row["block_location"]
    last_node_location=csv_row["last_node_location"]
    
    last_node_type=csv_row["last_node_type"]
    last_node_value=csv_row["last_node_value"]
    
        
    node_location_obj=NodeLocation.parse(filepath,node_location)
    block_location_obj=NodeLocation.parse(filepath,block_location)
    last_node_location_obj=NodeLocation.parse(filepath,last_node_location)
    
    
    last_cfg_node_id=CfgNode.id_construct(last_node_type,last_node_value,last_node_location_obj)
    
    return CfgNode(ast_node_id=AstNode.id_construct(node_type,node_value,node_location_obj),
                   owner_block_id=BasicBlockInfo.id_construct(last_cfg_node_id,block_location_obj),
                   location=node_location_obj)

def cfgedge_create(csv_row:dict):
    start_type=csv_row["start_type"]
    start_value=csv_row["start_value"]
    end_type=csv_row["end_type"]
    end_value=csv_row["end_value"]
    
    file_path=csv_row["filepath"]
    start_location=csv_row["start_location"]
    end_location=csv_row["end_location"]
    
    start_location_obj=NodeLocation.parse(file_path,start_location)
    end_location_obj=NodeLocation.parse(file_path,end_location)
    
    start_ast_id=AstNode.id_construct(start_type,start_value,start_location_obj)
    end_ast_id=AstNode.id_construct(end_type,end_value,end_location_obj)
    
    return CfgEdge(start_ast_id,end_ast_id,file_path)

def basicblockinfo_create(csv_row:dict):
    scope=csv_row["scope"]
    filepath=csv_row["filepath"]
    location=csv_row["block_location"]
    last_node_type=csv_row["node_type"]
    last_node_value=csv_row["node_value"]
    last_node_location=csv_row["node_location"]
    
    scope=Scope.parse(scope)
    block_location=NodeLocation.parse(filepath,location)
    node_location_obj=NodeLocation.parse(filepath,last_node_location)


    
    return BasicBlockInfo(scope=scope,
                          location=block_location,
                          last_node_id=CfgNode.id_construct(last_node_type,last_node_value,node_location_obj))

def basicblockedge_create(csv_row:dict):
    scope=csv_row["scope"]
    filepath=csv_row["filepath"]
    
    start_location=csv_row["start_location"]
    end_location=csv_row["end_location"]
    
    start_node_type=csv_row["start_node_type"]
    start_node_value=csv_row["start_node_value"]
    start_node_location=csv_row["start_node_location"]

    end_node_type=csv_row["end_node_type"]
    end_node_value=csv_row["end_node_value"]
    end_node_location=csv_row["end_node_location"]
    
    start_node_location_obj=NodeLocation.parse(filepath,start_node_location)
    end_node_location_obj=NodeLocation.parse(filepath,end_node_location)
    start_location_obj=NodeLocation.parse(filepath,start_location)
    end_location_obj=NodeLocation.parse(filepath,end_location)
    
    start_node_id=CfgNode.id_construct(start_node_type,start_node_value,start_node_location_obj)
    end_node_id=CfgNode.id_construct(end_node_type,end_node_value,end_node_location_obj)
    
    start_block_id=BasicBlockInfo.id_construct(start_node_id,start_location_obj)
    end_block_id=BasicBlockInfo.id_construct(end_node_id,end_location_obj)
    
    return BasicBlockEdge(start_block_id,end_block_id,filepath)

def ast_generate_from_csv(ast_node_path,ast_edge_path)->tuple[list[AstNode],list[AstEdge]]:
    nodes=[]
    edges=[]
    
    with open(ast_node_path,'r') as f:
        reader=csv.DictReader(f)
        #columns=reader.fieldnames

        for row in reader:
            nodes.append(astnode_create(row))
    
    with open(ast_edge_path,'r') as f:
        reader=csv.DictReader(f)        
    
        for row in reader:
            edges.append(astedge_create(row))
    
    return nodes,edges

def ast_tree_create(ast_node_path,ast_edge_path):
    ast_nodes,ast_edges=ast_generate_from_csv(ast_node_path,ast_edge_path)
    filepath2node={}
    filepath2edge={}
    
    for e in ast_nodes:
        filepath=e.location.file_path
        if filepath not in filepath2node:
            filepath2node[filepath]=[]
        filepath2node[filepath].append(e)
        
    for e in ast_edges:
        filepath=e.file_path
        if filepath not in filepath2edge:
            filepath2edge[filepath]=[]
        filepath2edge[filepath].append(e)
    
    common_filepaths=set(filepath2edge).union(filepath2node)
        
    tree_create_handle=lambda x:AstTree(filepath2node.get(x,[]),filepath2edge.get(x,[]))
    return list(map(tree_create_handle,common_filepaths))

def basicblock_info_generate_from_csv(basicblock_info_path,basicblock_connect_path):
    nodes=[]
    edges=[]
    
    with open(basicblock_info_path,'r') as f:
        reader=csv.DictReader(f)
        #columns=reader.fieldnames

        for row in reader:
            nodes.append(basicblockinfo_create(row))
    
    with open(basicblock_connect_path,'r') as f:
        reader=csv.DictReader(f)        
    
        for row in reader:
            edges.append(basicblockedge_create(row))
    
    return nodes,edges

def basicblock_graph_create(basicblock_info_path,basicblock_connect_path):
    nodes,edges=basicblock_info_generate_from_csv(basicblock_info_path,basicblock_connect_path)
    filepath2node={}
    filepath2edge={}
    
    for e in nodes:
        filepath=e.block_location.file_path
        if filepath not in filepath2node:
            filepath2node[filepath]=[]
        filepath2node[filepath].append(e)
        
    for e in edges:
        filepath=e.file_path
        if filepath not in filepath2edge:
            filepath2edge[filepath]=[]
        filepath2edge[filepath].append(e)
    
    common_filepaths=set(filepath2edge).union(filepath2node)
        
    basicblock_create_handle=lambda x:BasicBlockGraph(filepath2node.get(x,[]),filepath2edge.get(x,[]))
    return list(map(basicblock_create_handle,common_filepaths))

def tokenize_python(codeql_src_zip_path,source_path):
    def token_node_create(token_info:tokenize.TokenInfo):
        return TokenNode(string=token_info.string,
                         type=token_info.type,
                         start=token_info.start,
                         end=token_info.end)
    #test_path='mnt/XiaoweiGuo/vscode/CodeCompletion/data/py150ksource/test/1/9/batch_1.py'
    if source_path[0] in ['\\','/']:
        source_path=source_path[1:]
    with zipfile.ZipFile(codeql_src_zip_path,'r') as zip_file:
        #print(zip_file.infolist())
        with zip_file.open(source_path,'r') as f:
            return list(map(token_node_create,tokenize.tokenize(f.readline)))
    
def source_token_generate_python(codeql_src_zip_path,source_path):
    return TokenSequence(tokenize_python(codeql_src_zip_path,source_path))


def source_generate_from_csv(query_dir,codeql_db_path):
    
    ast_edge_path=os.path.join(query_dir,'ast-edge.csv')
    ast_node_path=os.path.join(query_dir,'ast-node.csv')
    
    basic_block_node_path=os.path.join(query_dir,'basic-block.csv')
    basic_block_edge_path=os.path.join(query_dir,'basic-block-connect.csv')
    codeql_src_zip_path=os.path.join(codeql_db_path,'src.zip')
    
    ast_trees=ast_tree_create(ast_node_path,ast_edge_path)
    basic_block_group=basicblock_graph_create(basic_block_node_path,basic_block_edge_path)
    token_seq_create_handle=lambda x:source_token_generate_python(codeql_src_zip_path,x.file_path)
    
    token_seqs=map(token_seq_create_handle,ast_trees)
    
    for i,e in enumerate(token_seqs):
        yield Source(ast_tree=ast_trees[i],basic_blocks=basic_block_group[i],token_seq=e)
        
